#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
智能负载均衡器
动态分配任务到工作线程，优化处理性能
"""

import os
import time
import threading
import multiprocessing
from queue import Queue, PriorityQueue, Empty
from typing import Dict, List, Tuple, Any, Optional, Callable
from dataclasses import dataclass, field
from enum import Enum
import psutil

from ..utils_module.logger import ObfuscationLogger


class WorkerStatus(Enum):
    """工作线程状态"""
    IDLE = "idle"
    BUSY = "busy"
    BLOCKED = "blocked"
    ERROR = "error"
    TERMINATED = "terminated"


@dataclass
class Task:
    """任务定义"""
    task_id: str
    file_path: str
    priority: int  # 优先级（数字越小优先级越高）
    estimated_time: float  # 预估处理时间（秒）
    dependencies: List[str] = field(default_factory=list)
    metadata: Dict[str, Any] = field(default_factory=dict)

    def __lt__(self, other):
        """比较优先级（用于优先队列）"""
        return self.priority < other.priority


@dataclass
class Worker:
    """工作线程信息"""
    worker_id: str
    thread: threading.Thread = None
    status: WorkerStatus = WorkerStatus.IDLE
    current_task: Optional[Task] = None
    tasks_completed: int = 0
    total_time: float = 0.0
    error_count: int = 0
    last_active: float = field(default_factory=time.time)


@dataclass
class LoadMetrics:
    """负载指标"""
    cpu_percent: float
    memory_percent: float
    active_workers: int
    queue_size: int
    avg_task_time: float
    throughput: float  # 任务/秒


class LoadBalancer:
    """智能负载均衡器"""

    def __init__(self, logger: ObfuscationLogger = None,
                 max_workers: int = None,
                 adaptive: bool = True):
        """
        初始化负载均衡器

        Args:
            logger: 日志记录器
            max_workers: 最大工作线程数（None表示自动）
            adaptive: 是否启用自适应调整
        """
        self.logger = logger

        # 工作线程配置
        if max_workers is None:
            # 自动确定：CPU核心数的80%，至少2个，最多16个
            cpu_count = multiprocessing.cpu_count()
            max_workers = max(2, min(16, int(cpu_count * 0.8)))

        self.max_workers = max_workers
        self.adaptive = adaptive

        # 任务队列
        self.task_queue = PriorityQueue()
        self.pending_tasks: Dict[str, Task] = {}
        self.completed_tasks: Dict[str, Any] = {}
        self.failed_tasks: Dict[str, str] = {}

        # 工作线程池
        self.workers: Dict[str, Worker] = {}
        self.worker_lock = threading.RLock()

        # 性能指标
        self.start_time = time.time()
        self.total_tasks = 0
        self.completed_count = 0
        self.failed_count = 0

        # 负载监控
        self.monitor_thread = None
        self.monitoring = False
        self.load_history: List[LoadMetrics] = []

        # 任务处理器
        self.task_processor: Optional[Callable] = None

        # 控制信号
        self.shutdown_event = threading.Event()
        self.running = False

    def distribute_tasks(self, tasks: List[Any], worker_count: int) -> List[List[Any]]:
        """分配任务给工作线程"""
        if not tasks:
            return []

        # 简单轮询分配
        distributed = [[] for _ in range(worker_count)]
        for i, task in enumerate(tasks):
            distributed[i % worker_count].append(task)

        return distributed

    def balance_tasks(self, tasks: List[Any]) -> Dict[str, List[Any]]:
        """平衡任务分配"""
        return {"balanced_tasks": tasks}

    def start_monitoring(self):
        """开始监控"""
        if not hasattr(self, 'monitoring') or self.monitoring:
            return

        self.monitoring = True
        self.monitor_thread = threading.Thread(
            target=self._monitor_loop,
            daemon=True
        )
        self.monitor_thread.start()

    def _monitor_loop(self):
        """监控循环"""
        while getattr(self, 'monitoring', False):
            try:
                # 简单的监控逻辑
                if hasattr(self, 'workers'):
                    active_workers = sum(1 for w in self.workers.values()
                                       if hasattr(w, 'status') and w.status == WorkerStatus.BUSY)

                    # 记录负载指标
                    current_time = time.time()
                    load_metric = type('LoadMetrics', (), {
                        'timestamp': current_time,
                        'active_workers': active_workers,
                        'queue_size': self.task_queue.qsize() if hasattr(self.task_queue, 'qsize') else 0
                    })()

                    if hasattr(self, 'load_history'):
                        self.load_history.append(load_metric)
                        # 保持历史记录在合理范围内
                        if len(self.load_history) > 100:
                            self.load_history.pop(0)

                time.sleep(1)  # 每秒检查一次

            except Exception as e:
                if hasattr(self, 'logger') and self.logger:
                    self.logger.log_error(f"负载监控错误: {e}")

    def add_task(self, task: Any, priority: int = 5) -> str:
        """添加任务到队列"""
        if not hasattr(self, 'pending_tasks'):
            self.pending_tasks = {}

        task_id = f"task_{int(time.time() * 1000)}"
        self.pending_tasks[task_id] = task

        # 简单地将任务添加到队列
        if hasattr(self, 'task_queue'):
            try:
                self.task_queue.put((priority, task_id, task))
            except:
                # 如果队列满了，直接处理
                self._process_task_directly(task)

        return task_id

    def _process_task_directly(self, task: Any):
        """直接处理任务"""
        try:
            if hasattr(self, 'task_processor') and self.task_processor:
                self.task_processor(task)
            else:
                # 默认处理逻辑
                time.sleep(0.001)  # 模拟处理时间
        except Exception as e:
            if hasattr(self, 'logger') and self.logger:
                self.logger.log_error(f"任务处理错误: {e}")

    def wait_for_completion(self, timeout: float = 30):
        """等待所有任务完成"""
        start_time = time.time()
        while time.time() - start_time < timeout:
            if (hasattr(self, 'task_queue') and
                hasattr(self.task_queue, 'empty') and
                self.task_queue.empty()):
                break
            time.sleep(0.1)

    def set_task_processor(self, processor: Callable[[Task], Any]):
        """
        设置任务处理器

        Args:
            processor: 任务处理函数
        """
        self.task_processor = processor

    def add_task(self, task: Task) -> bool:
        """
        添加任务到队列

        Args:
            task: 任务对象

        Returns:
            是否成功添加
        """
        if task.task_id in self.pending_tasks:
            if self.logger:
                self.logger.log_warning(f"任务 {task.task_id} 已存在")
            return False

        # 检查依赖
        for dep in task.dependencies:
            if dep not in self.completed_tasks:
                # 依赖未完成，暂存
                self.pending_tasks[task.task_id] = task
                return True

        # 添加到队列
        self.task_queue.put(task)
        self.total_tasks += 1

        return True

    def add_batch_tasks(self, tasks: List[Task]):
        """
        批量添加任务

        Args:
            tasks: 任务列表
        """
        added = 0
        for task in tasks:
            if self.add_task(task):
                added += 1

        if self.logger:
            self.logger.log_operation(
                "批量任务",
                f"添加了 {added}/{len(tasks)} 个任务"
            )

    def start(self):
        """启动负载均衡器"""
        if self.monitoring:
            return

        self.monitoring = True
        self.shutdown_event.clear()

        # 启动初始工作线程
        initial_workers = min(4, self.max_workers)
        for i in range(initial_workers):
            self._spawn_worker(f"worker_{i}")

        # 启动监控线程
        self.monitor_thread = threading.Thread(
            target=self._monitor_loop,
            daemon=True
        )
        self.monitor_thread.start()

        if self.logger:
            self.logger.log_operation(
                "负载均衡器",
                f"已启动 {initial_workers} 个工作线程"
            )

    def stop(self, timeout: float = 30):
        """
        停止负载均衡器

        Args:
            timeout: 超时时间（秒）
        """
        self.monitoring = False
        self.shutdown_event.set()

        # 等待所有任务完成
        deadline = time.time() + timeout

        while time.time() < deadline:
            if self.task_queue.empty() and all(
                w.status == WorkerStatus.IDLE
                for w in self.workers.values()
            ):
                break
            time.sleep(0.1)

        # 终止所有工作线程
        with self.worker_lock:
            for worker in self.workers.values():
                worker.status = WorkerStatus.TERMINATED

        # 等待线程结束
        for worker in self.workers.values():
            if worker.thread and worker.thread.is_alive():
                worker.thread.join(timeout=1)

        if self.logger:
            self.logger.log_operation(
                "负载均衡器",
                f"已停止: 完成={self.completed_count}, 失败={self.failed_count}"
            )

    def _spawn_worker(self, worker_id: str) -> Worker:
        """
        创建新的工作线程

        Args:
            worker_id: 工作线程ID

        Returns:
            工作线程对象
        """
        with self.worker_lock:
            if worker_id in self.workers:
                return self.workers[worker_id]

            worker = Worker(worker_id=worker_id)
            worker.thread = threading.Thread(
                target=self._worker_loop,
                args=(worker,),
                daemon=True
            )
            worker.thread.start()

            self.workers[worker_id] = worker
            return worker

    def _worker_loop(self, worker: Worker):
        """
        工作线程主循环

        Args:
            worker: 工作线程对象
        """
        while not self.shutdown_event.is_set():
            if worker.status == WorkerStatus.TERMINATED:
                break

            try:
                # 获取任务（超时1秒）
                task = self.task_queue.get(timeout=1)

                # 更新状态
                worker.status = WorkerStatus.BUSY
                worker.current_task = task
                worker.last_active = time.time()

                # 处理任务
                start_time = time.time()

                if self.task_processor:
                    result = self.task_processor(task)

                    # 记录完成
                    self.completed_tasks[task.task_id] = result
                    self.completed_count += 1

                    # 更新统计
                    elapsed = time.time() - start_time
                    worker.tasks_completed += 1
                    worker.total_time += elapsed

                    # 检查依赖任务
                    self._check_pending_tasks(task.task_id)

                    if self.logger:
                        self.logger.log_operation(
                            f"任务完成",
                            f"{task.task_id} 耗时 {elapsed:.2f}秒"
                        )
                else:
                    # 无处理器，任务失败
                    self.failed_tasks[task.task_id] = "无任务处理器"
                    self.failed_count += 1

            except Empty:
                # 队列为空
                worker.status = WorkerStatus.IDLE
                worker.current_task = None

            except Exception as e:
                # 任务处理失败
                worker.error_count += 1

                if worker.current_task:
                    self.failed_tasks[worker.current_task.task_id] = str(e)
                    self.failed_count += 1

                    if self.logger:
                        self.logger.log_error(
                            f"任务失败: {worker.current_task.task_id} - {e}"
                        )

                worker.status = WorkerStatus.ERROR
                worker.current_task = None

                # 错误恢复
                time.sleep(1)
                worker.status = WorkerStatus.IDLE

    def _check_pending_tasks(self, completed_task_id: str):
        """
        检查并释放依赖已完成的任务

        Args:
            completed_task_id: 完成的任务ID
        """
        ready_tasks = []

        for task_id, task in list(self.pending_tasks.items()):
            if completed_task_id in task.dependencies:
                # 移除已完成的依赖
                task.dependencies.remove(completed_task_id)

                # 如果所有依赖都已完成
                if not task.dependencies:
                    ready_tasks.append(task)
                    del self.pending_tasks[task_id]

        # 将就绪任务加入队列
        for task in ready_tasks:
            self.task_queue.put(task)

    def _monitor_loop(self):
        """监控循环"""
        while self.monitoring:
            # 收集负载指标
            metrics = self._collect_metrics()
            self.load_history.append(metrics)

            # 保持历史记录在合理范围内
            if len(self.load_history) > 100:
                self.load_history = self.load_history[-50:]

            # 自适应调整
            if self.adaptive:
                self._adaptive_scaling(metrics)

            # 健康检查
            self._health_check()

            # 监控间隔
            time.sleep(5)

    def _collect_metrics(self) -> LoadMetrics:
        """
        收集负载指标

        Returns:
            负载指标
        """
        # CPU和内存使用率
        cpu_percent = psutil.cpu_percent(interval=0.1)
        memory_percent = psutil.virtual_memory().percent

        # 活跃工作线程数
        active_workers = sum(
            1 for w in self.workers.values()
            if w.status == WorkerStatus.BUSY
        )

        # 队列大小
        queue_size = self.task_queue.qsize()

        # 平均任务时间
        total_time = sum(w.total_time for w in self.workers.values())
        total_completed = sum(w.tasks_completed for w in self.workers.values())
        avg_task_time = total_time / total_completed if total_completed > 0 else 0

        # 吞吐量
        elapsed = time.time() - self.start_time
        throughput = self.completed_count / elapsed if elapsed > 0 else 0

        return LoadMetrics(
            cpu_percent=cpu_percent,
            memory_percent=memory_percent,
            active_workers=active_workers,
            queue_size=queue_size,
            avg_task_time=avg_task_time,
            throughput=throughput
        )

    def _adaptive_scaling(self, metrics: LoadMetrics):
        """
        自适应扩缩容

        Args:
            metrics: 当前负载指标
        """
        current_workers = len(self.workers)

        # 扩容条件
        if (metrics.queue_size > current_workers * 10 and
            metrics.cpu_percent < 80 and
            metrics.memory_percent < 80 and
            current_workers < self.max_workers):

            # 增加工作线程
            new_worker_count = min(
                current_workers + 2,
                self.max_workers
            )

            for i in range(current_workers, new_worker_count):
                self._spawn_worker(f"worker_{i}")

            if self.logger:
                self.logger.log_operation(
                    "自适应扩容",
                    f"工作线程从 {current_workers} 增加到 {new_worker_count}"
                )

        # 缩容条件
        elif (metrics.queue_size < current_workers * 2 and
              current_workers > 2 and
              metrics.active_workers < current_workers * 0.3):

            # 减少工作线程
            idle_workers = [
                w for w in self.workers.values()
                if w.status == WorkerStatus.IDLE
            ]

            # 终止部分空闲线程
            to_terminate = min(2, len(idle_workers))
            for i in range(to_terminate):
                idle_workers[i].status = WorkerStatus.TERMINATED

            if self.logger and to_terminate > 0:
                self.logger.log_operation(
                    "自适应缩容",
                    f"终止了 {to_terminate} 个空闲工作线程"
                )

    def _health_check(self):
        """健康检查"""
        current_time = time.time()

        with self.worker_lock:
            for worker in list(self.workers.values()):
                # 检查死锁或卡住的线程
                if (worker.status == WorkerStatus.BUSY and
                    current_time - worker.last_active > 300):  # 5分钟

                    if self.logger:
                        self.logger.log_warning(
                            f"工作线程 {worker.worker_id} 可能卡住"
                        )

                    # 标记为错误
                    worker.status = WorkerStatus.ERROR

                    # 重新调度任务
                    if worker.current_task:
                        self.task_queue.put(worker.current_task)
                        worker.current_task = None

                # 移除已终止的线程
                if (worker.status == WorkerStatus.TERMINATED and
                    worker.thread and not worker.thread.is_alive()):
                    del self.workers[worker.worker_id]

    def get_statistics(self) -> Dict[str, Any]:
        """
        获取统计信息

        Returns:
            统计信息字典
        """
        # 当前指标
        current_metrics = self._collect_metrics()

        # 工作线程统计
        worker_stats = {
            'total': len(self.workers),
            'busy': sum(1 for w in self.workers.values()
                       if w.status == WorkerStatus.BUSY),
            'idle': sum(1 for w in self.workers.values()
                       if w.status == WorkerStatus.IDLE),
            'error': sum(1 for w in self.workers.values()
                        if w.status == WorkerStatus.ERROR)
        }

        # 任务统计
        task_stats = {
            'total': self.total_tasks,
            'completed': self.completed_count,
            'failed': self.failed_count,
            'pending': len(self.pending_tasks),
            'queued': self.task_queue.qsize()
        }

        # 性能统计
        performance_stats = {
            'avg_task_time': current_metrics.avg_task_time,
            'throughput': current_metrics.throughput,
            'cpu_usage': current_metrics.cpu_percent,
            'memory_usage': current_metrics.memory_percent
        }

        return {
            'workers': worker_stats,
            'tasks': task_stats,
            'performance': performance_stats
        }

    def wait_for_completion(self, timeout: Optional[float] = None) -> bool:
        """
        等待所有任务完成

        Args:
            timeout: 超时时间（秒）

        Returns:
            是否在超时前完成
        """
        start_time = time.time()

        while True:
            # 检查是否完成
            if (self.task_queue.empty() and
                len(self.pending_tasks) == 0 and
                all(w.status == WorkerStatus.IDLE
                    for w in self.workers.values())):
                return True

            # 检查超时
            if timeout and time.time() - start_time > timeout:
                return False

            time.sleep(0.1)

    def get_task_result(self, task_id: str) -> Optional[Any]:
        """
        获取任务结果

        Args:
            task_id: 任务ID

        Returns:
            任务结果，如果未完成返回None
        """
        return self.completed_tasks.get(task_id)

    def get_task_error(self, task_id: str) -> Optional[str]:
        """
        获取任务错误信息

        Args:
            task_id: 任务ID

        Returns:
            错误信息，如果未失败返回None
        """
        return self.failed_tasks.get(task_id)

    def estimate_completion_time(self) -> float:
        """
        估算完成时间

        Returns:
            预计剩余时间（秒）
        """
        remaining = self.task_queue.qsize() + len(self.pending_tasks)

        if remaining == 0:
            return 0

        # 基于历史吞吐量估算
        if self.load_history:
            avg_throughput = sum(m.throughput for m in self.load_history[-10:]) / min(10, len(self.load_history))
            if avg_throughput > 0:
                return remaining / avg_throughput

        # 基于当前速度估算
        elapsed = time.time() - self.start_time
        if self.completed_count > 0 and elapsed > 0:
            rate = self.completed_count / elapsed
            return remaining / rate

        # 无法估算
        return -1